<?php

namespace plugin\ai\app\controller;

use support\Request;
use support\Response;
use Throwable;
use Workerman\Connection\AsyncTcpConnection;
use Workerman\Protocols\Http\Chunk;

/**
 * 消息模块
 */
class MessageController
{

    /**
     * 发消息
     * @param Request $request
     * @return Response
     */
    public function send(Request $request): ?Response
    {
        // 提示词
        $prompt = $request->post('prompt');
        $model = $request->post('model');
        if ($model === 'dall.e') {
            return $this->draw($request, $prompt);
        }
        return $this->chat($request, $prompt);
    }

    /**
     * 聊天
     * @param Request $request
     * @param $prompt
     * @return Response
     */
    protected function chat(Request $request, $prompt): Response
    {
        // chatgpt 配置
        $gptConfig = config('plugin.ai.chatgpt');
        $chatGptApiHost = $gptConfig['api_host'];
        $apiKey = $request->post('apiKey') ?: $gptConfig['api_key'] ;
        $model = $request->post('model') ?: $gptConfig['model'];
        $maxTokens = (int)$request->post('maxTokens') ?: $gptConfig['max_tokens'];
        $temperature = (float)($request->post('temperature') ?: $gptConfig['temperature']);
        $lastMessages = $request->post('context', []);
        $rolePrompt = $request->post('rolePrompt', '');
        $messages = [];
        // 上下文有限制，不能太长
        $maxContextLength = $model === 'gpt-3.5-turbo-16k' ? 24000 : 6000;
        $contextLength = strlen($rolePrompt) + strlen($prompt) + $maxTokens * 1.5;
        if ($model !== 'gpt-3.5-turbo-16k') {
            $maxTokens = (int)min($maxTokens, ($maxTokens * 1.5 - strlen($rolePrompt) - strlen($prompt)) / 1.5 - 10);
        }
        foreach (array_reverse($lastMessages) as $item) {
            if (!isset($item['role']) || !isset($item['content'])) {
                continue;
            }
            $contextLength = $contextLength + strlen($item['content']);
            if ($contextLength < $maxContextLength) {
                $messages[] = $item;
            }
        }
        // 增加角色预设指令
        $messages[] = [
            'role' => 'assistant',
            'content' => $rolePrompt,
        ];
        $messages = array_reverse($messages);
        $messages[] = ["role" => "user", "content" => $prompt];

        foreach($messages as $key => $message) {
            if (!isset($message['content']) || $message['content'] === '') {
                unset($messages[$key]);
            }
        }
        $messages = array_values($messages);

        // 设置API URL和参数
        $data = array(
            'model' => $model,
            'max_tokens' => $maxTokens,
            'temperature' => $temperature,
            'stream' => true,
            'messages' => $messages
        );

        // 向 chatgpt api 发送数据
        [$schema, $host] = explode('://', $chatGptApiHost, 2);
        $con = new AsyncTcpConnection("tcp://$host", ['ssl' => [
            'verify_peer' => false,
        ]]);
        $con->transport = $schema === 'https' ? 'ssl' : 'tcp';
        $body = json_encode($data);
        $bodyLength = strlen($body);
        $con->send(
            "POST /v1/chat/completions HTTP/1.1\r\n" .
            "Host: $host\r\n" .
            "Connection: close\r\n" .
            "Content-Type: application/json\r\n" .
            "Authorization: Bearer $apiKey\r\n" .
            "Content-Length: $bodyLength\r\n\r\n" .
            $body
        );
        // 获取浏览器链接
        $connection = $request->connection;
        // 失败时
        $con->onError = function ($con, $code, $msg) use ($connection) {
            $connection->send(new Chunk(json_encode(['error' => ['message' => $msg]])));
            $connection->send(new Chunk(''));
        };
        // api接口返回数据时
        $con->onMessage = function ($con, $buffer) use ($connection) {
            static $headerCompleted = false, $header = '';
            if (!$headerCompleted) {
                $header .= $buffer;
                if (!$position = strpos($header, "\r\n\r\n")) {
                    return;
                }
                $headerCompleted = true;
                $bodyLength = 0;
                if (preg_match("/Content-Length: (\d+)\r\n/", $header, $match)) {
                    $bodyLength = $match[1];
                }
                if(!$buffer = substr($header, $position + 4)) {
                    return;
                }
                if ($bodyLength) {
                    $connection->send(new Chunk($buffer));
                    $connection->send(new Chunk(''));
                    return;
                }
            }
            $connection->send($buffer, true);
        };
        $con->connect();
        // 向浏览器发送头部响应
        return response("\n")->withHeaders([
            "Content-Type" => "application/octet-stream",
            "Transfer-Encoding" => "chunked",
        ]);
    }

    /**
     * 画图
     * @param Request $request
     * @param $prompt
     * @return Response
     */
    protected function draw(Request $request, $prompt): Response
    {
        // chatgpt 配置
        $gptConfig = config('plugin.ai.chatgpt');
        $chatGptApiHost = $gptConfig['api_host'];
        $apiKey = $request->post('apiKey') ?: $gptConfig['api_key'] ;

        // 设置API URL和参数
        $data = array(
            'prompt' => $prompt,
            'n' => 1,
            'size' => "512x512"
        );

        // 向 chatgpt api 发送数据
        [$schema, $host] = explode('://', $chatGptApiHost, 2);
        $con = new AsyncTcpConnection("tcp://$host", ['ssl' => [
            'verify_peer' => false,
        ]]);
        $con->transport = $schema === 'https' ? 'ssl' : 'tcp';
        $body = json_encode($data);
        $bodyLength = strlen($body);
        $con->send(
            "POST /v1/images/generations HTTP/1.1\r\n" .
            "Host: $host\r\n" .
            "Connection: close\r\n" .
            "Content-Type: application/json\r\n" .
            "Authorization: Bearer $apiKey\r\n" .
            "Content-Length: $bodyLength\r\n\r\n" .
            $body
        );
        // 获取浏览器链接
        $connection = $request->connection;
        // 失败时
        $con->onError = function ($con, $code, $msg) use ($connection) {
            $connection->send(new Chunk(json_encode(['error' => ['message' => $msg]])));
            $connection->send(new Chunk(''));
        };
        // api接口返回数据时
        $con->onMessage = function ($con, $buffer) use ($connection) {
            static $headerCompleted = false, $header = '';
            if (!$headerCompleted) {
                $header .= $buffer;
                if (!$position = strpos($header, "\r\n\r\n")) {
                    return;
                }
                $headerCompleted = true;
                if(!$buffer = substr($header, $position + 4)) {
                    return;
                }
                $connection->send(new Chunk($buffer));
                $connection->send(new Chunk(''));
            }
            $connection->send(new Chunk($buffer));
            $connection->send(new Chunk(''));
        };
        $con->connect();
        // 向浏览器发送头部响应
        return response("\n")->withHeaders([
            "Content-Type" => "application/octet-stream",
            "Transfer-Encoding" => "chunked",
        ]);
    }

}